#!/usr/bin/env python3

import numpy as np
import torch
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import json
import os
from utils_cylinder import create_cylinder_drum, locate_cylinder_particles


class SimpleCylinderSimulator:
    
    def __init__(self, config):
        self.config = config
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"Device: {self.device}")
        
        self.gravity = -9.81
        self.dt = 0.001
        self.friction_coeff = 0.3
        self.restitution = 0.7
        
    def generate_particles(self):
        domain = self.config["simulation_domain"]
        n_particles = self.config["particles"]["n_kinematic_particles"]
        inner_radius = self.config["cylinder_info"]["inner_radius"]
        
        particles = []
        velocities = []
        
        for _ in range(n_particles):
            r = np.random.uniform(0, inner_radius * 0.8)
            theta = np.random.uniform(0, 2*np.pi)
            x = np.random.uniform(-0.02, 0.02)
            y = r * np.cos(theta)
            z = r * np.sin(theta)
            
            particles.append([x, y, z])
            
            init_vel = self.config["particles"]["initial_velocity"]
            vx = init_vel[0] + np.random.normal(0, 0.01)
            vy = np.random.normal(0, 0.01)
            vz = np.random.normal(0, 0.01)
            velocities.append([vx, vy, vz])
        
        return (torch.tensor(particles, dtype=torch.float32, device=self.device),
                torch.tensor(velocities, dtype=torch.float32, device=self.device))
    
    def check_cylinder_collision(self, positions, velocities, cylinder_center, rotation_angle):
        inner_radius = self.config["cylinder_info"]["inner_radius"]
        outer_radius = self.config["cylinder_info"]["outer_radius"]
        length_x = self.config["cylinder_info"]["length_x"]
        
        rel_pos = positions - cylinder_center
        
        cos_angle = torch.cos(rotation_angle)
        sin_angle = torch.sin(rotation_angle)
        
        y_rot = rel_pos[:, 1] * cos_angle + rel_pos[:, 2] * sin_angle
        z_rot = -rel_pos[:, 1] * sin_angle + rel_pos[:, 2] * cos_angle
        
        r = torch.sqrt(y_rot**2 + z_rot**2)
        
        inner_collision = r < inner_radius
        outer_collision = r > outer_radius
        
        x_left_collision = rel_pos[:, 0] < 0
        x_right_collision = rel_pos[:, 0] > length_x
        
        new_velocities = velocities.clone()
        new_positions = positions.clone()
        
        if torch.any(inner_collision):
            mask = inner_collision
            r_unit_y = y_rot[mask] / (r[mask] + 1e-8)
            r_unit_z = z_rot[mask] / (r[mask] + 1e-8)
            
            v_radial = (new_velocities[mask, 1] * r_unit_y + 
                       new_velocities[mask, 2] * r_unit_z)
            
            new_velocities[mask, 1] -= 2 * v_radial * r_unit_y * self.restitution
            new_velocities[mask, 2] -= 2 * v_radial * r_unit_z * self.restitution
            
            new_positions[mask, 1] = cylinder_center[1] + r_unit_y * inner_radius * 1.01
            new_positions[mask, 2] = cylinder_center[2] + r_unit_z * inner_radius * 1.01
        
        if torch.any(x_left_collision):
            mask = x_left_collision
            new_velocities[mask, 0] = torch.abs(new_velocities[mask, 0]) * self.restitution
            new_positions[mask, 0] = cylinder_center[0] + 0.001
        
        if torch.any(x_right_collision):
            mask = x_right_collision
            new_velocities[mask, 0] = -torch.abs(new_velocities[mask, 0]) * self.restitution
            new_positions[mask, 0] = cylinder_center[0] + length_x - 0.001
        
        return new_positions, new_velocities
    
    def simulate_step(self, positions, velocities, cylinder_center, rotation_angle):
        velocities[:, 1] += self.gravity * self.dt
        new_positions = positions + velocities * self.dt
        
        new_positions, new_velocities = self.check_cylinder_collision(
            new_positions, velocities, cylinder_center, rotation_angle
        )
        
        drag_coeff = 0.01
        new_velocities *= (1 - drag_coeff * self.dt)
        
        return new_positions, new_velocities
    
    def run_simulation(self, n_steps, cylinder_center, rotation_angle):
        positions, velocities = self.generate_particles()
        trajectory = [positions.clone()]
        
        print(f"Running simulation: {n_steps} steps...")
        for step in range(n_steps):
            positions, velocities = self.simulate_step(
                positions, velocities, cylinder_center, rotation_angle
            )
            
            if step % (n_steps // 10) == 0:
                print(f"Step {step}/{n_steps}")
                trajectory.append(positions.clone())
        
        trajectory.append(positions.clone())
        return torch.stack(trajectory)


def compute_performance_metrics(trajectory):
    final_positions = trajectory[-1]
    initial_positions = trajectory[0]
    
    centroid_initial = torch.mean(initial_positions, dim=0)
    centroid_final = torch.mean(final_positions, dim=0)
    centroid_displacement = torch.norm(centroid_final - centroid_initial)
    
    final_distances = torch.norm(final_positions - centroid_final.unsqueeze(0), dim=1)
    dispersion = torch.std(final_distances)
    
    x_displacement = torch.mean(final_positions[:, 0] - initial_positions[:, 0])
    
    return {
        'centroid_displacement': centroid_displacement.item(),
        'dispersion': dispersion.item(),
        'x_displacement': x_displacement.item(),
        'final_centroid': centroid_final.cpu().numpy(),
        'initial_centroid': centroid_initial.cpu().numpy()
    }


def optimization_loop(simulator, config):
    cylinder_center = torch.tensor(config["cylinder_info"]["center_position"], 
                                 dtype=torch.float32, device=simulator.device)
    
    rotation_angles = torch.linspace(0, np.pi, 10)
    target_x_displacement = 0.02
    
    best_angle = 0
    best_loss = float('inf')
    results = []
    
    print("Starting angle optimization...")
    
    for i, angle in enumerate(rotation_angles):
        print(f"\nTesting angle {i+1}/10: {angle.item():.3f} rad ({angle.item()*180/np.pi:.1f}°)")
        
        trajectory = simulator.run_simulation(
            n_steps=100, 
            cylinder_center=cylinder_center, 
            rotation_angle=angle
        )
        
        metrics = compute_performance_metrics(trajectory)
        loss = (metrics['x_displacement'] - target_x_displacement)**2
        
        results.append({
            'angle': angle.item(),
            'loss': loss,
            'metrics': metrics,
            'trajectory': trajectory
        })
        
        if loss < best_loss:
            best_loss = loss
            best_angle = angle.item()
        
        print(f"  X displacement: {metrics['x_displacement']:.4f}m")
        print(f"  Centroid displacement: {metrics['centroid_displacement']:.4f}m") 
        print(f"  Dispersion: {metrics['dispersion']:.4f}m")
        print(f"  Loss: {loss:.6f}")
    
    return best_angle, best_loss, results


def create_visualizations(results, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    
    angles = [r['angle'] for r in results]
    losses = [r['loss'] for r in results]
    x_displacements = [r['metrics']['x_displacement'] for r in results]
    
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))
    
    ax1.plot(np.array(angles) * 180 / np.pi, losses, 'bo-')
    ax1.set_xlabel('Rotation Angle (deg)')
    ax1.set_ylabel('Loss')
    ax1.set_title('Optimization - Loss vs Angle')
    ax1.grid(True)
    
    ax2.plot(np.array(angles) * 180 / np.pi, x_displacements, 'ro-')
    ax2.axhline(y=0.02, color='g', linestyle='--', label='Target')
    ax2.set_xlabel('Rotation Angle (deg)')
    ax2.set_ylabel('X Displacement (m)')
    ax2.set_title('X Displacement vs Angle')
    ax2.legend()
    ax2.grid(True)
    
    plt.tight_layout()
    plt.savefig(f'{output_dir}/optimization_results.png', dpi=150)
    plt.close()
    
    best_result = min(results, key=lambda x: x['loss'])
    trajectory = best_result['trajectory']
    
    fig = plt.figure(figsize=(15, 10))
    
    ax1 = fig.add_subplot(2, 3, 1, projection='3d')
    initial_pos = trajectory[0].cpu().numpy()
    final_pos = trajectory[-1].cpu().numpy()
    
    ax1.scatter(initial_pos[:, 0], initial_pos[:, 1], initial_pos[:, 2], 
               c='blue', s=10, alpha=0.6, label='Initial')
    ax1.scatter(final_pos[:, 0], final_pos[:, 1], final_pos[:, 2], 
               c='red', s=10, alpha=0.6, label='Final')
    ax1.set_xlabel('X')
    ax1.set_ylabel('Y')
    ax1.set_zlabel('Z')
    ax1.set_title('Particle Trajectories')
    ax1.legend()
    
    ax2 = fig.add_subplot(2, 3, 2)
    time_steps = len(trajectory)
    centroids_x = [torch.mean(traj[:, 0]).item() for traj in trajectory]
    centroids_y = [torch.mean(traj[:, 1]).item() for traj in trajectory]
    
    ax2.plot(range(time_steps), centroids_x, 'b-', label='X centroid')
    ax2.plot(range(time_steps), centroids_y, 'r-', label='Y centroid')
    ax2.set_xlabel('Time step')
    ax2.set_ylabel('Centroid position (m)')
    ax2.set_title('Centroid evolution')
    ax2.legend()
    ax2.grid(True)
    
    ax3 = fig.add_subplot(2, 3, 3)
    ax3.scatter(initial_pos[:, 0], initial_pos[:, 1], c='blue', s=10, alpha=0.6, label='Initial')
    ax3.scatter(final_pos[:, 0], final_pos[:, 1], c='red', s=10, alpha=0.6, label='Final')
    ax3.set_xlabel('X')
    ax3.set_ylabel('Y')
    ax3.set_title('X-Y projection')
    ax3.legend()
    ax3.set_aspect('equal')
    
    ax4 = fig.add_subplot(2, 3, 4)
    ax4.scatter(initial_pos[:, 1], initial_pos[:, 2], c='blue', s=10, alpha=0.6, label='Initial')
    ax4.scatter(final_pos[:, 1], final_pos[:, 2], c='red', s=10, alpha=0.6, label='Final')
    
    theta = np.linspace(0, 2*np.pi, 100)
    inner_r = 0.05
    ax4.plot(inner_r * np.cos(theta), inner_r * np.sin(theta), 'k-', linewidth=2, label='Wall')
    
    ax4.set_xlabel('Y')
    ax4.set_ylabel('Z')
    ax4.set_title('Y-Z projection')
    ax4.legend()
    ax4.set_aspect('equal')
    
    ax5 = fig.add_subplot(2, 3, 5)
    x_displacements_all = [r['metrics']['x_displacement'] for r in results]
    dispersions = [r['metrics']['dispersion'] for r in results]
    
    ax5.scatter(x_displacements_all, dispersions, c=angles, cmap='viridis', s=50)
    ax5.set_xlabel('X displacement (m)')
    ax5.set_ylabel('Dispersion (m)')
    ax5.set_title('Displacement vs Dispersion')
    cbar = plt.colorbar(ax5.collections[0], ax=ax5)
    cbar.set_label('Angle (rad)')
    
    ax6 = fig.add_subplot(2, 3, 6)
    metrics_names = ['Centroid\nDisp.', 'X\nDisp.', 'Dispersion']
    best_metrics = best_result['metrics']
    values = [best_metrics['centroid_displacement'], 
              best_metrics['x_displacement'], 
              best_metrics['dispersion']]
    
    bars = ax6.bar(metrics_names, values, color=['blue', 'red', 'green'])
    ax6.set_ylabel('Value (m)')
    ax6.set_title(f'Best metrics\n(angle: {best_result["angle"]*180/np.pi:.1f}°)')
    
    for bar, value in zip(bars, values):
        height = bar.get_height()
        ax6.text(bar.get_x() + bar.get_width()/2., height + 0.001,
                f'{value:.4f}', ha='center', va='bottom')
    
    plt.tight_layout()
    plt.savefig(f'{output_dir}/detailed_analysis.png', dpi=150, bbox_inches='tight')
    plt.close()
    
    print(f"\nVisualization saved to: {output_dir}")


def main():
    print("="*60)
    print("Cylinder simulation demo")
    print("="*60)
    
    config_path = 'inverse_barrier/config_cylinder.json'
    with open(config_path, 'r') as f:
        config = json.load(f)
    
    print("Config:")
    print(f"  Inner diameter: {config['cylinder_info']['inner_radius']*2000:.0f}mm")
    print(f"  Outer diameter: {config['cylinder_info']['outer_radius']*2000:.0f}mm")
    print(f"  Length: {config['cylinder_info']['length_x']*1000:.0f}mm")
    print(f"  Particles: {config['particles']['n_kinematic_particles']}")
    
    simulator = SimpleCylinderSimulator(config)
    
    print("\nGenerating cylinder geometry...")
    cylinder_particles = create_cylinder_drum(
        inner_radius=config["cylinder_info"]["inner_radius"],
        outer_radius=config["cylinder_info"]["outer_radius"],
        length_x=config["cylinder_info"]["length_x"],
        spacing=config["simulation_domain"]["particle_spacing"]
    )
    print(f"Wall particles: {len(cylinder_particles)}")
    
    best_angle, best_loss, results = optimization_loop(simulator, config)
    
    output_dir = config["save_options"]["output_dir"]
    create_visualizations(results, output_dir)
    
    final_results = {
        'best_rotation_angle_rad': best_angle,
        'best_rotation_angle_deg': best_angle * 180 / np.pi,
        'best_loss': best_loss,
        'all_results': results,
        'config': config,
        'n_cylinder_particles': len(cylinder_particles)
    }
    
    torch.save(final_results, f'{output_dir}/simulation_results.pt')
    
    print("\n" + "="*60)
    print("Simulation complete!")
    print(f"Best rotation angle: {best_angle*180/np.pi:.1f}° ({best_angle:.3f} rad)")
    print(f"Minimum loss: {best_loss:.6f}")
    print(f"Results saved in: {output_dir}")
    print("Generated files:")
    print("  - optimization_results.png")
    print("  - detailed_analysis.png")
    print("  - simulation_results.pt")
    print("="*60)


if __name__ == "__main__":
    main() 